import os
import operator
import asyncio
from dotenv import load_dotenv
from typing import Annotated, Any, TypedDict, List

from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI

from langgraph.graph import StateGraph, END

# --- 1. 环境准备 ---
load_dotenv()
# 增加一个检查，确保环境变量被加载
if not os.getenv("OPENAI_BASE_URL") or not os.getenv("OPENAI_API_KEY"):
    print("错误：请确保在 .env 文件中或环境变量中设置了 OPENAI_BASE_URL 和 OPENAI_API_KEY")
    exit()

# 为了演示，我们使用一个简单的搜索工具
try:
    from langchain_community.tools.tavily_search import TavilySearchResults
    search_tool = TavilySearchResults(max_results=2)
except ImportError:
    from langchain_community.tools import DuckDuckGoSearchRun
    search_tool = DuckDuckGoSearchRun()

# --- 2. 状态定义 ---
# 状态需要跟踪分发的任务、并行收集的结果以及最终报告
class AgentState(TypedDict):
    company: str
    research_tasks: List[str]
    research_results: Annotated[list, operator.add]
    # 新增一个字段，专门用于存放聚合后的简报
    briefing: str
    report: str

# --- 3. 代理与节点定义 (增加调试信息) ---

# a. 通用 LLM
llm = ChatOpenAI(model="qwen-plus-latest", base_url=os.getenv("OPENAI_BASE_URL"))

# b. 分发节点 (Distributor)
def distributor_node(state: AgentState):
    print("\n--- [DEBUG] 进入分发节点 (distributor_node) ---")
    company = state["company"]
    tasks = [
        f"分析 {company} 的最新财务状况和股票表现。",
        f"总结关于 {company} 的最新新闻和公众舆论。",
        f"调查 {company} 的核心技术和最新产品动态。",
    ]
    print(f"--- [DEBUG] 已创建任务: {tasks} ---")
    return {"research_tasks": tasks}

# c. 并行研究节点 (Parallel Researchers) - 增加错误处理
research_prompt = PromptTemplate.from_template(
    """你是一位专业的研究分析师。你的任务是根据以下要求进行深入研究，并提供一份简洁的总结报告。

研究要求: {task}

你的研究报告:"""
)

# 为单个研究任务创建一个带错误处理的 chain
async def run_single_research_task(task: str) -> AIMessage:
    """包装单个研究任务，增加打印和错误捕获"""
    print(f"--- [DEBUG] 开始并行研究任务: '{task[:30]}...' ---")
    try:
        # 原始的 chain 逻辑
        chain = research_prompt | llm
        result = await chain.ainvoke({"task": task})
        print(f"--- [DEBUG] 完成并行研究任务: '{task[:30]}...' ---")
        return result
    except Exception as e:
        print(f"--- [ERROR] 研究任务 '{task[:30]}...' 失败: {e} ---")
        # 即使失败，也返回一个消息，避免整个流程崩溃
        return AIMessage(content=f"研究任务失败，错误信息: {e}")

# 使用 RunnableLambda 将我们的 async 函数包装成一个可运行的组件
# research_chain 将接收一个任务字符串，并返回一份研究报告
research_chain_runnable = RunnableLambda(run_single_research_task)

# 这个函数用于将并行节点的输出（一个AIMessage列表）包装成 langgraph 期望的字典格式
def format_research_output(results: List[AIMessage]) -> dict:
    """这个函数保持不变，它只做格式化"""
    content_results = [result.content for result in results]
    print(f"--- [DEBUG] 格式化研究成果，收到 {len(content_results)} 条。---")
    return {"research_results": content_results}

# .map() 是 LCEL 中实现并行处理的核心
# 我们将带错误处理的 runnable 应用于 map, 然后再连接格式化函数
research_node = research_chain_runnable.map() | RunnableLambda(format_research_output)

# d. 聚合节点 (Aggregator)
def aggregator_node(state: AgentState):
    print("\n--- [DEBUG] 进入聚合节点 (aggregator_node) ---")
    research_results = state["research_results"]
    
    briefing = "\n\n---\n\n".join(research_results)
    print("--- [DEBUG] 已将研究成果聚合成简报 ---")
    return {"briefing": briefing}

# e. 报告撰写节点 (Writer)
writer_prompt = PromptTemplate.from_template(
    """你是一位资深的行业分析顾问。请根据以下研究简报，为 {company} 撰写一份全面的市场研究报告。
报告应该结构清晰，包含引言、各维度分析和总结。

研究简报:
{briefing}

你的最终市场研究报告:"""
)
writer_chain = writer_prompt | llm

async def writer_node(state: AgentState):
    print("\n--- [DEBUG] 进入报告撰写节点 (writer_node) ---")
    company = state["company"]
    briefing = state["briefing"]
    
    try:
        # 注意：这个 ainvoke 调用应该被 astream_events 拦截以实现流式输出。
        # 正如你所发现的，目前它似乎在等待完整结果返回后才继续，
        # 这很可能是因为底层模型或API端点不支持与LangChain兼容的流式响应。
        report = await writer_chain.ainvoke({
            "company": company,
            "briefing": briefing,
        })
        print("--- [DEBUG] 最终报告已生成 ---")
        # 我们在这里返回完整报告。流式处理的逻辑在 astream_events 循环中，负责监听事件。
        return {"report": report.content}
    except Exception as e:
        print(f"--- [ERROR] 报告撰写失败: {e} ---")
        return {"report": f"报告生成失败，错误信息: {e}"}


# --- 4. 构建图 ---
workflow = StateGraph(AgentState)

workflow.add_node("distributor", distributor_node)
workflow.add_node("researcher", research_node)
workflow.add_node("aggregator", aggregator_node)
workflow.add_node("writer", writer_node)

workflow.set_entry_point("distributor")

workflow.add_edge("distributor", "researcher")
workflow.add_edge("researcher", "aggregator")
workflow.add_edge("aggregator", "writer")

workflow.set_finish_point("writer")


# --- 5. 编译与运行 (流式版本) ---
app = workflow.compile()

async def main():
    initial_state = {
        "company": "特斯拉（Tesla）",
        "research_results": [], # 初始化为空列表
    }
    
    print("---开始执行, 请确保 .env 文件已正确配置---")
    final_state = None
    # 使用 astream_events 来实时观察图的执行流程
    async for event in app.astream_events(initial_state, version="v2"):
        kind = event["event"]
        name = event["name"]
        
        # 实时打印 writer 节点的流式输出
        if kind == "on_chat_model_stream" and name == "writer":
             content = event["data"]["chunk"].content
             if content:
                 print(content, end="", flush=True)

        # 捕获最终状态
        if kind == "on_graph_end":
            final_state = event["data"]["output"]
            
    print("\n\n---流程结束---")
    if final_state:
        # 最后的打印：
        # - 如果流式输出工作，这里是最终结果的确认。
        # - 如果流式输出不工作，这里将是首次看到完整报告的地方。
        print("\n最终市场研究报告:")
        print(final_state.get("report", "报告生成失败，未在最终状态中找到。"))

if __name__ == "__main__":
    asyncio.run(main()) 